# conding=utf-8
import logging
import struct
from operator import attrgetter
from ryu.base import app_manager
from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, DEAD_DISPATCHER
from ryu.controller.handler import CONFIG_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.ofproto import ofproto_v1_0
from ryu.ofproto import ether
from ryu.lib.packet import packet
from ryu.lib.packet import ethernet
from ryu.lib.packet import ipv4
from ryu.lib.packet import arp
from ryu.lib.ip     import ipv4_to_bin
from ryu.lib        import addrconv

import ryu.networkx as nx

import network_aware
PREAMBLE = 'aa:aa:aa:aa:aa:aa'


def convert_ip_to_int(ip):
    '''
    From string to int
    10.0.0.2 -> 167772162
    '''
    tmp = 0
    for b in ipv4_to_bin(ip):
        tmp |= ord(b)
        tmp = tmp << 8
    tmp = tmp >> 8
    return tmp

class Shortest_Route(app_manager.RyuApp):
    OFP_VERSIONS = [ofproto_v1_0.OFP_VERSION]
    _CONTEXTS = {
        "Network_Aware": network_aware.Network_Aware,
    }

    def __init__(self, *args, **kwargs):
        super(Shortest_Route, self).__init__(*args, **kwargs)
        self.network_aware = kwargs["Network_Aware"]
        
        self.datapaths = self.network_aware.dps

        # links   :(src_dpid,dst_dpid)->(src_port,dst_port)
        self.link_to_port = self.network_aware.link_to_port

        # {(dpid,port)->host_ip}
        self.access_table = self.network_aware.access_table

        # Tor switch
        self.edge_sw = self.network_aware.edge_sw

        # Core switch
        self.core_sw = self.network_aware.core_sw

        self.graph = self.network_aware.whole_graph

        # Multicast Address ---> {dpid1:[port1,port2,port3],...,dpid n:[port1,port2,..])}
        # The first dpid1,port1 means the sender hosts, and the other means the recivers
        self.multicast_groups = self.network_aware.multicast_groups

    @set_ev_cls(ofp_event.EventOFPStateChange,
                [MAIN_DISPATCHER, DEAD_DISPATCHER])
    def _state_change_handler(self, ev):
        datapath = ev.datapath
        if ev.state == MAIN_DISPATCHER:
            if not datapath.id in self.datapaths:
                self.logger.debug('register datapath: %016x', datapath.id)
                self.datapaths[datapath.id] = datapath
        elif ev.state == DEAD_DISPATCHER:
            if datapath.id in self.datapaths:
                self.logger.debug('unregister datapath: %016x', datapath.id)
                del self.datapaths[datapath.id]

    def add_flow(self, datapath, p, match, actions, idle_timeout=0, hard_timeout=0):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        mod = parser.OFPFlowMod(
            datapath=datapath, match=match, cookie=0,
            command=ofproto.OFPFC_ADD, idle_timeout=idle_timeout, hard_timeout=hard_timeout,
            priority=p,
            flags=ofproto.OFPFF_SEND_FLOW_REM, actions=actions)
        
        datapath.send_msg(mod)

    def install_multicast_flow(self, multicast_tree, multicast_group, multicast_group_addr):
        
        quene = []
        inport_quene = []
        quene.append(multicast_group.keys()[0])
        inport_quene.append(multicast_group.values()[0][0])

        while len(quene) > 0:
            datapath = self.datapaths[quene[0]]                    
            ofproto = datapath.ofproto
            parser = datapath.ofproto_parser

            #outport
            #link ---> quene[0]---successor[x]
            actions = []
            out_port = ofproto.OFPP_LOCAL

            successor = multicast_tree.edge[quene[0]].keys()
            for x in xrange(0,len(successor)):
                if len(multicast_tree.edge[successor[x]].keys()) > 0 or successor[x] in multicast_group.keys(): # the switch without hosts and successors should not be included
                    quene.append(successor[x])
                    inport_quene.append(self.link_to_port[(quene[0],successor[x])][1])
                    out_port = self.link_to_port[(quene[0], successor[x])][0]
                    self.logger.info("outport!: %s", out_port)
                    actions.append(parser.OFPActionOutput(out_port))

            if quene[0] in multicast_group.keys():
                if quene[0] != multicast_group.keys()[0]:
                    for x in xrange(0,len(multicast_group[quene[0]])):
                        out_port = multicast_group[quene[0]][x]
                        self.logger.info("outport!!: %s", out_port)
                        actions.append(parser.OFPActionOutput(out_port))
                else:
                    for x in xrange(1,len(multicast_group[quene[0]])):
                        out_port = multicast_group[quene[0]][x]
                        self.logger.info("outport!!!: %s", out_port)
                        actions.append(parser.OFPActionOutput(out_port))

            match = parser.OFPMatch(
                in_port = inport_quene[0],
                nw_dst = convert_ip_to_int(multicast_group_addr))
            self.add_flow(
                datapath, 1, match, actions,
                idle_timeout=20, hard_timeout=40)

            self.logger.info("switch: %s, in_port: %s, nw_dst %s", quene[0], inport_quene[0], multicast_group_addr)
            
            del quene[0]
            del inport_quene[0]

    def install_flow(self, path, flow_info, buffer_id, data):
        '''
            path=[dpid1, dpid2, dpid3...]
            flow_info=(eth_type, src_ip, dst_ip, in_port)
        '''
        # first flow entry
        in_port = flow_info[3]
        assert path
        datapath_first = self.datapaths[path[0]]
        ofproto = datapath_first.ofproto
        parser = datapath_first.ofproto_parser
        out_port = ofproto.OFPP_LOCAL

        # inter_link
        if len(path) > 2:
            self.logger.info("path > 2")
            for i in xrange(1, len(path) - 1):
                port = self.get_link2port(path[i - 1], path[i])
                port_next = self.get_link2port(path[i], path[i + 1])
                if port:
                    src_port, dst_port = port[1], port_next[0]
                    datapath = self.datapaths[path[i]]
                    ofproto = datapath.ofproto
                    parser = datapath.ofproto_parser
                    actions = []

                    actions = [parser.OFPActionOutput(dst_port)]
                    match = parser.OFPMatch(
                        in_port=src_port,
                        dl_type=flow_info[0],
                        nw_src=convert_ip_to_int(flow_info[1]),
                        nw_dst=convert_ip_to_int(flow_info[2]))
                    self.add_flow(
                        datapath, 1, match, actions,
                        idle_timeout=10, hard_timeout=30)

                    # inter links pkt_out
                    msg_data = None
                    if buffer_id == ofproto.OFP_NO_BUFFER:
                        msg_data = data

                    out = parser.OFPPacketOut(
                        datapath=datapath, buffer_id=buffer_id,
                        data=msg_data, in_port=src_port, actions=actions)

                    datapath.send_msg(out)

        if len(path) > 1:
            self.logger.info("path > 1")
            # the  first flow entry
            port_pair = self.get_link2port(path[0], path[1])
            out_port = port_pair[0]

            actions = []
            actions = [parser.OFPActionOutput(out_port)]
            match = parser.OFPMatch(
                in_port=in_port,
                dl_type=flow_info[0],
                nw_src=convert_ip_to_int(flow_info[1]),
                nw_dst=convert_ip_to_int(flow_info[2]))
            self.add_flow(datapath_first,
                          1, match, actions, idle_timeout=10, hard_timeout=30)

            # the last hop: tor -> host
            datapath = self.datapaths[path[-1]]
            ofproto = datapath.ofproto
            parser = datapath.ofproto_parser
            actions = []
            src_port = self.get_link2port(path[-2], path[-1])[1]
            dst_port = None

            for key in self.access_table.keys():
                if flow_info[2] == self.access_table[key]:
                    dst_port = key[1]
                    break
            actions = [parser.OFPActionOutput(dst_port)]
            match = parser.OFPMatch(
                in_port=src_port,
                dl_type=flow_info[0],
                nw_src=convert_ip_to_int(flow_info[1]),
                nw_dst=convert_ip_to_int(flow_info[2]))

            self.add_flow(
                datapath, 1, match, actions, idle_timeout=10, hard_timeout=30)

            # first pkt_out
            actions = []

            actions = [parser.OFPActionOutput(out_port)]
            msg_data = None
            if buffer_id == ofproto.OFP_NO_BUFFER:
                msg_data = data

            out = parser.OFPPacketOut(
                datapath=datapath_first, buffer_id=buffer_id,
                data=msg_data, in_port=in_port, actions=actions)

            datapath_first.send_msg(out)

            # last pkt_out
            actions = []
            actions = [parser.OFPActionOutput(dst_port)]
            msg_data = None
            if buffer_id == ofproto.OFP_NO_BUFFER:
                msg_data = data

            out = parser.OFPPacketOut(
                datapath=datapath, buffer_id=buffer_id,
                data=msg_data, in_port=src_port, actions=actions)

            datapath.send_msg(out)

        else:  # src and dst on the same
            out_port = None
            actions = []
            for key in self.access_table.keys():
                if flow_info[2] == self.access_table[key]:
                    out_port = key[1]
                    break

            actions = [parser.OFPActionOutput(out_port)]

            match = parser.OFPMatch(
                in_port=in_port,
                dl_type=flow_info[0],
                nw_src=convert_ip_to_int(flow_info[1]),
                nw_dst=convert_ip_to_int(flow_info[2]))
            
            self.add_flow(
                datapath_first, 1, match, actions,
                idle_timeout=10, hard_timeout=30)

            # pkt_out
            msg_data = None
            if buffer_id == ofproto.OFP_NO_BUFFER:
                msg_data = data

            out = parser.OFPPacketOut(
                datapath=datapath_first, buffer_id=buffer_id,
                data=msg_data, in_port=in_port, actions=actions)

            datapath_first.send_msg(out)

    def get_host_location(self, host_ip):
        for key in self.access_table:
            if self.access_table[key] == host_ip:
                return key
        self.logger.debug("%s location is not found." % host_ip)
        return None

    def get_link2port(self, src_dpid, dst_dpid):
        if (src_dpid, dst_dpid) in self.link_to_port:
            return self.link_to_port[(src_dpid, dst_dpid)]
        else:
            self.logger.debug("Link to port is not found.")
            return None

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def _packet_in_handler(self, ev):
        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        in_port = msg.in_port
        pkt = packet.Packet(msg.data)

        preamble_bin = struct.unpack_from("!6sH", pkt.data)
        preamble = addrconv.mac.bin_to_text(preamble_bin[0])
        if preamble != PREAMBLE:
            dst_mac, src_mac, eth_type = struct.unpack_from("!6s6sH", pkt.data)
            arp_pkt = pkt.get_protocol(arp.arp)
            ip_pkt = pkt.get_protocol(ipv4.ipv4)

            if isinstance(arp_pkt, arp.arp):
                arp_src_ip = arp_pkt.src_ip
                arp_dst_ip = arp_pkt.dst_ip

                # multicast arp request
                if arp_dst_ip in self.multicast_groups:
                    self.logger.info("src_sw %s dst_sw %s", arp_src_ip, arp_dst_ip);
                    
                    #consturc multicast tree
                    reciver_table = self.multicast_groups[arp_dst_ip].keys()
                    if len(reciver_table) == 1: #local edge switch
                        multicast_sub = self.graph.subgraph(reciver_table)
                    else: #cross different edge switches
                        multicast_sub = self.graph.subgraph(self.core_sw + reciver_table)
                    multicast_tree = nx.bfs_tree(nx.minimum_spanning_tree(multicast_sub),reciver_table[0])

                    #install the multicast flows
                    self.install_multicast_flow(multicast_tree, self.multicast_groups[arp_dst_ip], arp_dst_ip)
                    
                    # construct the arp reply and send data back
                    actions = [parser.OFPActionOutput(in_port)]
                    ARP_Reply = packet.Packet()
                    ARP_Reply.add_protocol(ethernet.ethernet(
                        ethertype = ether.ETH_TYPE_ARP,
                        dst = addrconv.mac.bin_to_text(src_mac),
                        src = addrconv.mac.bin_to_text(src_mac)))
                    ARP_Reply.add_protocol(arp.arp(
                        opcode = arp.ARP_REPLY,
                        src_mac = addrconv.mac.bin_to_text(src_mac),
                        src_ip = arp_dst_ip,
                        dst_mac = addrconv.mac.bin_to_text(src_mac),
                        dst_ip = arp_src_ip))

                    ARP_Reply.serialize()

                    out = datapath.ofproto_parser.OFPPacketOut(
                        datapath = datapath,
                        buffer_id = datapath.ofproto.OFP_NO_BUFFER,
                        in_port = datapath.ofproto.OFPP_CONTROLLER,
                        actions = actions, data = ARP_Reply.data)
                    datapath.send_msg(out)

                    self.logger.info("send arp multicast reply back")

                # unicast arp request
                else:
                    result = self.get_host_location(arp_dst_ip)
                    if result:  # host record in access table.
                        datapath_dst, out_port = result[0], result[1]
                        actions = [parser.OFPActionOutput(out_port)]
                        datapath = self.datapaths[datapath_dst]

                        out = parser.OFPPacketOut(
                            datapath=datapath,
                            buffer_id=ofproto.OFP_NO_BUFFER,
                            in_port=ofproto.OFPP_CONTROLLER,
                            actions=actions, data=msg.data)
                        datapath.send_msg(out)
                        self.logger.info("arp packet out !!!!!!!!!!!!!!!!")

            if isinstance(ip_pkt, ipv4.ipv4):

                ip_src = ip_pkt.src
                ip_dst = ip_pkt.dst

                result = None
                src_sw = None
                dst_sw = None

                self.logger.info("src_sw %s dst_sw %s", src_sw, dst_sw);

                src_location = self.get_host_location(ip_src)
                dst_location = self.get_host_location(ip_dst)

                if src_location and dst_location:
                    src_sw = src_location[0]
                    dst_sw = dst_location[0]
                    self.logger.info("src_sw %s dst_sw %s", src_sw, dst_sw)

                    result = nx.shortest_path(self.graph, src_sw, dst_sw)

                
                if result:
                    self.logger.info(
                        " PATH[%s --> %s]:%s\n" % (ip_src, ip_dst, result))

                    flow_info = (eth_type, ip_src, ip_dst, in_port)
                    self.install_flow(result, flow_info, msg.buffer_id, msg.data)
                else:
                    # Reflesh the topology database.
                    self.network_aware.get_topology()